from fastapi import FastAPI, HTTPException, Request
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel, Field, ConfigDict
from typing import List, Dict, Optional, Literal, Union, Any
import logging
import os
import sys

# 设置环境变量
os.environ["NER_MODEL_PATH"] = "/appslog/NER/Medical-NER-ONNX"  # 本地ONNX模型路径
os.environ["FINANCE_NER_MODEL_PATH"] = "/appslog/NER/nlp_raner_named-entity-recognition_chinese-base-finance"  # 金融NER模型路径
os.environ["XINFERENCE_SERVER_URL"] = "http://10.128.10.186:9997"  # xinference服务地址
# 禁用符号链接警告
os.environ["HF_HUB_DISABLE_SYMLINKS_WARNING"] = "1"
# 禁用Hugging Face遥测
os.environ["HF_HUB_DISABLE_TELEMETRY"] = "1"
# 设置transformers缓存目录
os.environ["TRANSFORMERS_CACHE"] = "./cache"

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

# 打印环境信息
logger.info(f"Python版本: {sys.version}")
try:
    import transformers
    logger.info(f"Transformers版本: {transformers.__version__}")
except:
    logger.warning("无法导入transformers库")

# 在导入服务之前打印环境变量
logger.info(f"NER_MODEL_PATH: {os.environ.get('NER_MODEL_PATH')}")
logger.info(f"FINANCE_NER_MODEL_PATH: {os.environ.get('FINANCE_NER_MODEL_PATH')}")

# 导入服务
from services.ner_service import NERService
from services.finance_ner_service import FinanceNERService
from services.std_service import StdService
from services.abbr_service import AbbrService
from services.corr_service import CorrService
from services.gen_service import GenService

# 创建 FastAPI 应用
app = FastAPI(
    title="医疗NLP工具箱",
    description="提供医疗文本处理的API服务，包括命名实体识别、术语标准化、缩写扩展等功能",
    version="1.0.0"
)

# 配置跨域资源共享
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],    
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# 初始化各个服务
try:
    logger.info("正在初始化命名实体识别服务...")
    ner_service = NERService()  # 命名实体识别服务
    logger.info("NER服务初始化完成")
except Exception as e:
    logger.error(f"NER服务初始化失败: {str(e)}")
    print(f"错误: 无法加载NER模型: {str(e)}")
    sys.exit(1)  # 如果NER服务无法初始化，程序退出

try:
    logger.info("正在初始化金融命名实体识别服务...")
    finance_ner_service = FinanceNERService()  # 金融命名实体识别服务
    logger.info("金融NER服务初始化完成")
except Exception as e:
    logger.error(f"金融NER服务初始化失败: {str(e)}")
    print(f"错误: 无法加载金融NER模型: {str(e)}")
    finance_ner_service = None


try:
    logger.info("正在初始化术语标准化服务...")
    try:
        # 使用Xinference嵌入模型
        standardization_service = StdService(
            provider="xinference",
            model="bge-large",
            db_path="db/snomed_bge_m3.db",
            collection_name="concepts_only_name"
        )
        logger.info("术语标准化服务初始化完成")
    except Exception as e:
        logger.warning(f"术语标准化服务初始化失败: {str(e)}")
        standardization_service = None
except Exception as e:
    logger.warning(f"术语标准化服务外部初始化失败: {str(e)}")
    standardization_service = None
    
# 设置后端环境变量
os.environ["TRANSFORMERS_OFFLINE"] = "1"

# 这些服务不依赖外部模型下载，可以正常初始化
logger.info("正在初始化其他服务...")
abbr_service = AbbrService()  # 缩写扩展服务
gen_service = GenService()    # 文本生成服务
corr_service = CorrService()  # 拼写纠正服务
logger.info("所有服务初始化完成")

# 基础模型类
class BaseInputModel(BaseModel):
    """基础输入模型，包含所有模型共享的字段"""
    model_config = ConfigDict(arbitrary_types_allowed=True)
    
    llmOptions: Dict[str, str] = Field(
        default_factory=lambda: {
            "provider": "ollama",
            "model": "qwen2.5:7b"
        },
        description="大语言模型配置选项"
    )

class EmbeddingOptions(BaseModel):
    """向量数据库配置选项"""
    provider: Literal["huggingface", "openai", "bedrock", "xinference"] = Field(
        default="xinference",
        description="向量数据库提供商"
    )
    model: str = Field(
        default="bge-large",
        description="嵌入模型名称"
    )
    dbName: str = Field(
        default="snomed_bge_m3",
        description="向量数据库名称"
    )
    collectionName: str = Field(
        default="concepts_only_name",
        description="集合名称"
    )

class TextInput(BaseInputModel):
    """文本输入模型，用于标准化和命名实体识别"""
    text: str = Field(..., description="输入文本")
    options: Dict[str, bool] = Field(
        default_factory=dict,
        description="处理选项"
    )
    termTypes: Dict[str, bool] = Field(
        default_factory=dict,
        description="术语类型"
    )
    embeddingOptions: EmbeddingOptions = Field(
        default_factory=EmbeddingOptions,
        description="嵌入选项"
    )

# 金融向量搜索输入模型
class FinanceVectorInput(BaseInputModel):
    """金融向量搜索输入模型"""
    query: str = Field(..., description="查询文本")
    options: Dict[str, Any] = Field(
        default_factory=dict,
        description="搜索选项"
    )

class AbbrInput(BaseInputModel):
    """缩写扩展输入模型"""
    text: str = Field(..., description="输入文本")
    context: str = Field(
        default="",
        description="上下文信息"
    )
    method: Literal["simple_ollama", "query_db_llm_rerank", "llm_rank_query_db"] = Field(
        default="simple_ollama",
        description="处理方法"
    )
    embeddingOptions: Optional[EmbeddingOptions] = Field(
        default_factory=EmbeddingOptions,
        description="向量数据库配置选项"
    )

class ErrorOptions(BaseModel):
    """错误生成选项"""
    probability: float = Field(
        default=0.3,
        description="错误生成概率",
        ge=0.0,
        le=1.0
    )
    maxErrors: int = Field(
        default=5,
        description="最大错误数量",
        ge=1
    )
    keyboard: Literal["qwerty", "azerty"] = Field(
        default="qwerty",
        description="键盘布局"
    )

class CorrInput(BaseInputModel):
    """拼写纠正输入模型"""
    text: str = Field(..., description="输入文本")
    method: Literal["correct_spelling", "add_mistakes"] = Field(
        default="correct_spelling",
        description="处理方法"
    )
    errorOptions: ErrorOptions = Field(
        default_factory=ErrorOptions,
        description="错误生成选项"
    )

class PatientInfo(BaseModel):
    """患者信息模型"""
    name: str = Field(..., description="患者姓名")
    age: Optional[int] = Field(
        None,
        description="患者年龄",
        ge=0
    )
    gender: Optional[Literal["M", "F", "O"]] = Field(
        None,
        description="患者性别"
    )
    medicalHistory: Optional[str] = Field(
        None,
        description="既往病史"
    )

class GenInput(BaseInputModel):
    """医疗内容生成输入模型"""
    patient_info: PatientInfo = Field(..., description="患者信息")
    symptoms: List[str] = Field(..., description="症状列表")
    diagnosis: str = Field(
        default="",
        description="诊断结果"
    )
    treatment: str = Field(
        default="",
        description="治疗方案"
    )
    method: Literal["generate_medical_note", "generate_differential_diagnosis", "generate_treatment_plan"] = Field(
        default="generate_medical_note",
        description="生成方法"
    )

# API 端点：术语标准化
@app.post("/api/std")
async def standardization(input: TextInput):
    try:
        # 记录请求信息
        logger.info(f"Received request: text={input.text}, options={input.options}, embeddingOptions={input.embeddingOptions}")

        # 配置术语类型
        all_medical_terms = input.options.pop('allMedicalTerms', False)
        term_types = {'allMedicalTerms': all_medical_terms}

        # 进行命名实体识别
        ner_results = ner_service.process(input.text, input.options, term_types)

        # 初始化标准化服务
        standardization_service = StdService(
            provider=input.embeddingOptions.provider,
            model=input.embeddingOptions.model,
            db_path=f"db/{input.embeddingOptions.dbName}.db",
            collection_name=input.embeddingOptions.collectionName
        )

        # 获取识别到的实体
        entities = ner_results.get('entities', [])
        if not entities:
            return {"message": "No medical terms have been recognized", "standardized_terms": []}

        # 标准化每个实体
        standardized_results = []
        for entity in entities:
            std_result = standardization_service.search_similar_terms(entity['word'])
            standardized_results.append({
                "original_term": entity['word'],
                "entity_group": entity['entity_group'],
                "standardized_results": std_result
            })

        return {
            "message": f"{len(entities)} medical terms have been recognized and standardized",
            "standardized_terms": standardized_results
        }

    except Exception as e:
        logger.error(f"Error in standardization processing: {str(e)}")
        raise HTTPException(status_code=500, detail=str(e))

# API 端点：命名实体识别
@app.post("/api/ner")
async def ner(input: TextInput):
    """
    命名实体识别API
    """
    try:
        return ner_service.process(input.text, input.options, input.termTypes)
    except Exception as e:
        logger.error(f"NER处理错误: {str(e)}")
        raise HTTPException(status_code=500, detail=str(e))

# API 端点：金融向量搜索
@app.post("/api/finance-vector")
async def finance_vector_search(input: FinanceVectorInput):
    """
    金融向量搜索API
    """
    # 导入服务
    from services.finance_vector_service import FinanceVectorService
    
    try:
        # 初始化服务（这里使用默认配置）
        finance_vector_service = FinanceVectorService()
        return finance_vector_service.process(input.query, input.options)
    except Exception as e:
        logger.error(f"金融向量搜索错误: {str(e)}")
        raise HTTPException(status_code=500, detail=str(e))

@app.post("/api/finance-ner")
async def finance_ner(input: TextInput):
    """
    金融命名实体识别API
    """
    if finance_ner_service is None:
        raise HTTPException(status_code=503, detail="金融NER服务未初始化")
    try:
        return finance_ner_service.process(input.text, input.options, input.termTypes)
    except Exception as e:
        logger.error(f"金融NER处理错误: {str(e)}")
        raise HTTPException(status_code=500, detail=str(e))

# API 端点：拼写纠正
@app.post("/api/corr")
async def correct_notes(input: CorrInput):
    try:
        if input.method == "correct_spelling":  # 拼写纠正
            return corr_service.correct_spelling(input.text, input.llmOptions)
        elif input.method == "add_mistakes":  # 添加错误（测试用）
            return corr_service.add_mistakes(input.text, input.errorOptions)
        else:
            raise HTTPException(status_code=400, detail="Invalid method")
    except Exception as e:
        logger.error(f"Error in correction processing: {str(e)}")
        raise HTTPException(status_code=500, detail=str(e))

# API 端点：缩写扩展
@app.post("/api/abbr")
async def expand_abbreviations(input: AbbrInput):
    try:
        if input.method == "simple_ollama":  # 简单扩展
            output = abbr_service.simple_ollama_expansion(input.text, input.llmOptions)
            return {"input": input.text, "output": output}
        elif input.method == "query_db_llm_rerank":  # 数据库查询+重排序
            return abbr_service.query_db_llm_rerank(
                input.text, 
                input.context, 
                input.llmOptions,
                input.embeddingOptions
            )
        elif input.method == "llm_rank_query_db":  # LLM扩展+数据库标准化
            return abbr_service.llm_rank_query_db(
                input.text, 
                input.context, 
                input.llmOptions,
                input.embeddingOptions
            )
        else:
            raise HTTPException(status_code=400, detail="Invalid method")
    except Exception as e:
        logger.error(f"Error in abbreviation expansion: {str(e)}")
        raise HTTPException(status_code=500, detail=str(e))

# API 端点：医疗文本生成
@app.post("/api/gen")
async def generate_medical_content(input: GenInput):
    try:
        if input.method == "generate_medical_note":  # 生成病历
            return gen_service.generate_medical_note(
                input.patient_info,
                input.symptoms,
                input.diagnosis,
                input.treatment,
                input.llmOptions
            )
        elif input.method == "generate_differential_diagnosis":  # 生成鉴别诊断
            return gen_service.generate_differential_diagnosis(
                input.symptoms,
                input.llmOptions
            )
        elif input.method == "generate_treatment_plan":  # 生成治疗计划
            return gen_service.generate_treatment_plan(
                input.diagnosis,
                input.patient_info,
                input.llmOptions
            )
        else:
            raise HTTPException(status_code=400, detail="Invalid method")
    except Exception as e:
        logger.error(f"Error in medical content generation: {str(e)}")
        raise HTTPException(status_code=500, detail=str(e))

# 启动服务器
if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)
